# 导入必要的库
import torch  # PyTorch库，用于深度学习模型的训练和推理
from tqdm import tqdm  # tqdm库，用于显示进度条
from modelscope import AutoTokenizer, AutoModel  # modelscope库，用于自动加载预训练模型和分词器
from peft import PeftModel, PeftConfig  # peft库，用于进行prompt-tuning高效微调

# 定义模型和预训练模型的路径
#model_dir = "../../chatglm3-6b"  # 预训练模型的路径
model_dir = "C:\\Users\\16014\\.cache\\modelscope\\hub\\models\\ZhipuAI\\chatglm3-6b"
peft_model_id = "./lora_saver/lora_query_key_value"  # PEFT微调模型的路径

# 在不计算梯度的情况下进行模型加载和预处理
with torch.no_grad():
    # 从预训练模型路径自动加载分词器
    tokenizer = AutoTokenizer.from_pretrained(model_dir, trust_remote_code=True)
    # 从预训练模型路径自动加载模型，并转换为半精度浮点数格式（节省显存），然后移动到GPU上
    model = AutoModel.from_pretrained(model_dir, trust_remote_code=True).half().cuda()

# 使用PEFT微调模型对原始模型进行微调
model = PeftModel.from_pretrained(model, peft_model_id)
# 将模型设置为评估模式（关闭dropout、batchnorm等层的影响）
model.eval()

# 定义对话历史和查询
history = []  # 对话历史，初始为空列表
query = "你是谁"  # 查询语句，即用户的输入问题
role = "user"  # 角色，这里设置为"user"表示是用户的输入

# 使用分词器构建聊天输入
inputs = tokenizer.build_chat_input(query, history=history, role=role)
# 将输入数据移动到GPU上
inputs = inputs.to('cuda')
#inputs = inputs.to('cpu')


# 定义结束标记和生成参数
eos_token_id = [tokenizer.eos_token_id, tokenizer.get_command("<|user|>"),
                tokenizer.get_command("<|observation|>")]  # 结束标记，包括普通的结束标记和特殊角色的结束标记
gen_kwargs = {"max_length": 1200, "num_beams": 1, "do_sample": True, "top_p": 0.95,
              "temperature": 0.95}  # 生成参数，包括最大长度、集束搜索宽度、是否进行采样、top-p参数和温度参数

# 生成输出
outputs = model.generate(**inputs, **gen_kwargs, eos_token_id=eos_token_id)
# 对输出进行处理，去掉输入的部分并转换为列表形式
outputs = outputs.tolist()[0][len(inputs["input_ids"][0]):-1]

# 使用分词器解码输出
response = tokenizer.decode(outputs)
# 处理响应，包括去除一些特殊标记和更新对话历史
response, history = model.process_response(response, history)

# 打印响应
print(response)